import torch
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from components import EncoderLSTM, DecoderLSTM, ResidualStack
from enum import Enum, unique
import logging
from attrdict import AttrDict

from datasets import ScanPosVelDataset

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

fmt = logging.Formatter("[%(name)s] %(levelname)s %(message)s")
console_handler = logging.StreamHandler()
console_handler.setFormatter(fmt)
logger.addHandler(console_handler)


class EncoderTrainer(nn.Module):
    def __init__(self, obs_len=10, pred_len=10, hidden_dim = 128):
        super(EncoderTrainer, self).__init__()

        self.obs_len = obs_len
        self.pred_len = pred_len
        self.target_len = None
        self.hidden_dim = hidden_dim // 2
        
        # TODO: formalize the DecoderLSTM arguments passing
        # TODO: figure out how many convtranspose layer before decoder in SOGMP
        self.encoder = EncoderLSTM(1, self.hidden_dim)
        self.decoder = DecoderLSTM(1, self.hidden_dim*2)
        # self.output_layer = ResidualStack()
        
    def forward(self, hist_local_maps):
        # B,T,CHW
        outputs = []
        assert hist_local_maps.size(1) == self.target_len

        outs = self.encoder(hist_local_maps)
        b, _, h, w = outs.size()
        decoder_hidden = self.decoder.convlstm_.init_hidden(b, image_size=[h ,w])
        for _ in range(self.pred_len):
            output, hidden, cell = self.decoder(outs, *decoder_hidden)
            decoder_hidden = [hidden, cell]
            outputs.append(output)
        outputs = torch.stack(outputs, dim=1)

        return outputs
    

class ObsEncoderTrainer(EncoderTrainer):
    def __init__(self, obs_len=10, pred_len=10, hidden_dim=128):
        super(ObsEncoderTrainer, self).__init__(obs_len, pred_len, hidden_dim)
        self.target_len = obs_len


class PredEncoderTrainer(EncoderTrainer):
    def __init__(self, obs_len=10, pred_len=10, hidden_dim=128):
        super(PredEncoderTrainer, self).__init__(obs_len, pred_len, hidden_dim)
        self.target_len = pred_len

@unique
class TrainMode(Enum):
    TRAIN_OBS_ENC = 1
    TRAIN_PRED_ENC = 2
    TRAIN_DECODER = 3
    TRAIN_CLASSIFIER = 4
    TRAIN_MODALITY_SYNTHESIZER = 5
    EVAL = 6

    @classmethod
    def detail(cls, mode):
        if mode == cls.TRAIN_OBS_ENC:
            return "mode: TRAIN observation encoder"
        elif mode == cls.TRAIN_PRED_ENC:
            return "mode: TRAIN prediction encoder"
        elif mode == cls.TRAIN_DECODER:
            return "mode: TRAIN decoder"
        elif mode == cls.TRAIN_CLASSIFIER:
            return "mode: TRAIN modality classifier"
        elif mode == cls.TRAIN_MODALITY_SYNTHESIZER:
            return "mode: TRAIN regressor"
        elif mode == cls.EVAL:
            return "mode: EVAL"
        else:
            return "mode: UNKNOWN"

class PCCSDataset(ScanPosVelDataset):
    def __init__(self, 
                 dset_path, 
                 split, 
                 num_points=1080, 
                 image_size=64, 
                 obs_len=10,
                 pred_len=10):
        super(PCCSDataset, self).__init__(dset_path,
                                          split,
                                          num_points,
                                          image_size,
                                          obs_len,
                                          pred_len)
        self.obs_len = obs_len
        self.pred_len = pred_len
        self.split = split
        self.num_points = num_points
        self.image_size = image_size
        self.dset_path = dset_path

        self.obs_enc = ['test'] * int(1e5)
        self.pred_enc = ['test'] * int(1e5)
        self.cluster_results = ['test'] * int(1e5)
        self.classifier_gt = ['test'] * int(1e5)
        self.classifier_weights = ['test'] * int(1e5)
        self.current_mode = ['test'] * int(1e5)
    
    def set_mode(self, mode: TrainMode):
        self.current_mode = mode

    def log_mode(self):
        logger.info(TrainMode.detail(self.current_mode))

    def __len__(self):
        return super(PCCSDataset, self).__len__()
    
    def __getitem__(self, index):
        data = super(PCCSDataset, self).__getitem__(index)
        if self.current_mode == TrainMode.TRAIN_OBS_ENC:
            self.log_mode()
            data_obs = {k: data[k][:self.obs_len] for k in data}
            data_pred = {k: data[k][self.obs_len:] for k in data}
            return AttrDict({"observation": data_obs,
                             "prediction": data_pred})
        elif self.current_mode == TrainMode.TRAIN_PRED_ENC:
            self.log_mode()
            data_input = {k: data[k][self.obs_len:] for k in data}
            data_target = {k: data[k][self.obs_len:] for k in data}
            return AttrDict({"input": data_input,
                             "target": data_target})
        elif self.current_mode == TrainMode.TRAIN_DECODER:
            self.log_mode()
            # assert self.obs_enc is not None and self.pred_enc is not None
            data_target = {k: data[k][self.obs_len:] for k in data}
            return AttrDict({
                "obs_enc": self.obs_enc[index],
                "pred_enc": self.pred_enc[index],
                "target": data_target
            })
        elif self.current_mode == TrainMode.TRAIN_CLASSIFIER:
            self.log_mode()
            # assert self.obs_enc is not None and self.pred_enc is not None
            # assert self.cluster_results is not None
            return AttrDict({
                "obs_enc": self.obs_enc[index],
                "pred_enc": self.pred_enc[index],
                "cluster_idx": self.cluster_results[index]
            })
        elif self.current_mode == TrainMode.TRAIN_MODALITY_SYNTHESIZER:
            self.log_mode()
            # assert self.obs_enc is not None
            return AttrDict({
                "obs_enc": self.obs_enc[index],
                "prob": self.classifier_gt[index]
            })
        elif self.current_mode == TrainMode.EVAL:
            self.log_mode()
            # assert self.obs_enc is not None
            data_pred = {k: data[k][self.obs_len:] for k in data}
            return AttrDict({
                "obs_enc": self.obs_enc[index],
                "pred": data_pred
            })
        else:
            logger.warn("UNKNOWN mode")
            raise NotImplementedError
            return

if __name__ == "__main__":
    # test_tensor = torch.zeros(8, 10, 1, 64, 64)
    # enc_trainer = EncoderTrainer()
    # output = enc_trainer(test_tensor)
    # print(output.shape)
    test_set = PCCSDataset('/home/yflei/data/OGM-datasets/OGM-Turtlebot2/train', 'train')
    with open('d_set_test.txt', 'w') as f:
        for mode in TrainMode:
            test_set.set_mode(mode)
            f.write(str(test_set[0])+'\n'*5)